from app import db
from datetime import datetime
from pgvector.sqlalchemy import Vector

class KnowledgeBase(db.Model):
    __tablename__ = 'knowledge_base'
    __table_args__ = {'schema': 'public'}
    
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(100), nullable=False)
    description = db.Column(db.Text)
    embedding_model = db.Column(db.String(50), nullable=False)
    created_at = db.Column(db.DateTime(timezone=True), nullable=False, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime(timezone=True), nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 关联文档
    documents = db.relationship('Document', backref='knowledge_base', lazy='dynamic')

class Document(db.Model):
    __tablename__ = 'document'
    __table_args__ = {'schema': 'public'}
    
    id = db.Column(db.Integer, primary_key=True)
    knowledge_base_id = db.Column(db.Integer, db.ForeignKey('public.knowledge_base.id', ondelete='CASCADE'), nullable=False)
    title = db.Column(db.String(200), nullable=False)
    content = db.Column(db.Text, nullable=False)
    embedding = db.Column(Vector(1024), nullable=True)  # 存储向量数据
    doc_metadata = db.Column(db.JSON)  # 重命名为 doc_metadata
    created_at = db.Column(db.DateTime(timezone=True), nullable=False, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime(timezone=True), nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow) 